import cv2
import numpy as np
from numpy import ndarray
from moudels.utills import read_image, reload_image
from moudels.dispaly_image import display_images, display_images_with_attack
from moudels.watermarker import Watermarker
from moudels.attacker import Attacker
from moudels import ResultEvaluator



    
# 主函数
def main():
    original_image_path = "lena.bmp"  # 替换为你的原始图像路径
    watermark_path = "watermark.bmp"  # 替换为你的水印图像路径

    # 读取原始图像和水印图像
    original_image = read_image(original_image_path)
    watermark = read_image(watermark_path)

    # 确保载体图像大小为512×512
    if original_image.shape != (512, 512):
        raise ValueError("水印图像必须是64×64大小的灰度图像！")

    # 确保水印大小为64×64
    if watermark.shape != (64, 64):
        raise ValueError("水印图像必须是64×64大小的灰度图像！")

    #嵌入水印
    Watermarker.alpha = 0.1
    watermarked_image = Watermarker.embed(original_image, watermark)
    
    #攻击
    attack_type = Attacker.AttackType.ADD_GAUSSIAN_NOISE
    attacked_image = Attacker.attack(watermarked_image, attack_type)

    #提取
    extracted_watermark = Watermarker.extract(attacked_image, original_image)

    # 计算 PSNR
    psnr_value = ResultEvaluator.psnr(original_image, watermarked_image)
    nc_value = ResultEvaluator.nc(watermark, extracted_watermark)

    #展示结果
    display_images_with_attack(original_image, watermarked_image,attacked_image, attack_type, extracted_watermark, psnr_value, nc_value)


if __name__ == "__main__":
    main()